import torch
import numpy as np
import random
from model import WatermarkRemover
from torchvision.transforms import transforms
import cv2

fontFace_list = [
    cv2.FONT_ITALIC,
    cv2.FONT_HERSHEY_PLAIN,
    cv2.FONT_HERSHEY_SIMPLEX,
    cv2.FONT_HERSHEY_DUPLEX,
    cv2.FONT_HERSHEY_COMPLEX,
    cv2.FONT_HERSHEY_TRIPLEX,
    cv2.FONT_HERSHEY_SCRIPT_SIMPLEX,
    cv2.FONT_HERSHEY_SCRIPT_COMPLEX
]

def prediction(model:WatermarkRemover, img_path, save_path, device):
    model.eval()
    img = cv2.imread(img_path)
    # 生成水印
    mask = np.zeros(shape=img.shape, dtype=np.uint8)
    for i in range(random.randint(2, 12)):
        text = ''
        for _ in range(11):
            text += str(random.randint(0, 9))
        fontFace = fontFace_list[random.randint(0, len(fontFace_list)-1)]
        fontScale = random.random() * 1.5 + 0.5
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        thickness = random.randint(1, 3)
        x, y = int(mask.shape[1]*random.random()), int(mask.shape[0]*random.random())

        # text_size = cv2.getTextSize(text, fontFace, fontScale, thickness)
        # w, h = text_size[0]

        cv2.putText(mask, text, (x, y), fontFace, fontScale, color, thickness)
        # cv2.rectangle(img, (x, y), (x+w, y-h), (0, 0, 255), 1)
    # 合成水印图片
    t = random.randint(0, 9)
    img = cv2.addWeighted(img, 1, mask, (t+1)*0.1, 0)
    
    inp = transforms.ToTensor()(img) / 255.
    inp = inp.unsqueeze(0)
    
    with torch.no_grad():
        out = model(inp.to(device))

    out = torch.squeeze(out)
    out = out.permute(1,2,0)
    out *= 255.
    img2 = out.to('cpu')
    img2 = img2.numpy().astype(np.uint8)

    # 拼接输出和输入
    h, w = img.shape[:2]
    h2, w2 = img2.shape[:2]
    m1 = np.zeros((max(h, h2), max(w, w2), 3), dtype=np.uint8)
    m2 = np.zeros((max(h, h2), max(w, w2), 3), dtype=np.uint8)
    m1[:h, :w, :] = img
    m2[:h2, :w2, :] = img2
    cat_img = np.vstack((m1, m2))

    cv2.imwrite(save_path, cat_img)

